import cv2
import mujoco_py
import matplotlib.pyplot as plt

# load scene in MuJoCo
model = mujoco_py.load_model_from_path('jointExample.xml')
sim = mujoco_py.MjSim(model)
offscreen = mujoco_py.MjRenderContext(sim=sim, device_id=0, offscreen=True, opengl_backend='glfw')
viewer = mujoco_py.MjViewer(sim)
i = 0
while True:
    i += 1
    sim.step()
    viewer.render()
    if i % 100 == 0:
        offscreen.render(100, 100, camera_id=0)
        img = offscreen.read_pixels(100, 100, depth=False)
        img = img[..., ::-1]
        img = cv2.flip(img, 0)
        cv2.imwrite("cam/obs_" + str(i) + ".png", img)
    print('c')
